The R-squared calculation is incorrect. The denominator should use the mean of the observed values (y), not the predicted values (ŷ). The standard definition of R² is $1 - SS_{res}/SS_{tot}$, where $SS_{tot}$ is calculated relative to the mean of the observations.

    return 1 - sum((y[y_nan] .- ŷ[y_nan]).^2) / sum((y[y_nan] .- mean(y[y_nan])).^2)

This logic has two significant issues:

1. Shape Mismatch: falses(length(first(y))) creates a 1D array. If the targets are multi-dimensional (e.g., (time, batch)), the bitwise OR operation .|| will fail. Use size(first(y)) instead of length.
2. Incorrect Masking: Computing a single global is_no_nan mask by ORing all targets is problematic. If target A has a NaN at an index where target B is valid, the global mask will be true at that index. Consequently, the loss for target A will be computed using the NaN value, resulting in a NaN total loss. Masks should be computed and applied per-target.

The function now returns predictors and forcings separately, but the warning check at line 115 (visible in context) still references predictors_forcing. Since predictors_forcing is initialized as an empty array at line 89 and never populated in the new logic, this warning will be triggered on every call. The check should be updated to verify if both predictors and forcings are empty.

There is a typo in line 109: ŷ (y with combining circumflex) is used instead of the argument ŷ (U+0177) defined at line 105. While Julia normalizes identifiers to NFC, mixing these characters is confusing and can lead to issues in environments with different normalization rules. Additionally, the commented-out code should be removed.

                y_t = y[target]
                ŷ_t = ŷ[target]

The line ps = ps |> cfg.gdev is redundant. The parameters ps and state st are already moved to the device as part of the piped operation in line 19.

        ps, st = LuxCore.setup(Random.default_rng(), model) |> cfg.gdev

indeed, line 20 is not needed here.

we should do dev/ Array at the batch loader level. Up to this point data could still be lazy.

I think we can evaluate this still on the GPU side and just pipe the result of into the cfg.dev function.

the cfg.gdev(loader) is already moving the data into the gpu device, the second line should be an operation on the gpu side already, hence |> cfg.gdev should not be needed, in principle

doing y(target) was a intended to mirror AxisKeys syntax, although admittedly it would be better to do an independent test for that, and similarly for DD.

For now I've written the interface for loss functions to accept named tuples. So I may have to add a dispatch for callables.

These should work again.

I think is better to do cfg.gdev(loader), is just a wrapper at this point, I think. Then, when loop over, the (x,y) will be sent to the gpu. If it works now, we can come back to this later.

It works to do cfg.gdev(loader) except for dimensional data, it's still lazy and is not a simple array like the rest.

The fields gdev and cdev are untyped in the TrainConfig struct. In Julia, untyped fields lead to type instability, which can significantly degrade performance because the compiler cannot specialize functions using these fields. Since these are used frequently for device transfers during training, it is highly recommended to provide type annotations, such as Lux.AbstractDevice or using type parameters.

The change from ka(variable = k) to ka[k] will break if ka is a KeyedArray, as KeyedArray indexing typically requires dimension names or positional indices. While this might work if ka is now a NamedTuple due to changes in prepare_data, it makes the function less robust if it's still intended to handle KeyedArray inputs.

The construction of forcing_nt, targets_nt, and masks_nt using a loop and the NamedTuple constructor is inefficient. Since the inputs are already NamedTuples (as per the changes in prepare_data), you can use map which is more idiomatic and performant in Julia. Additionally, calling Array(v) on every batch is redundant if the data is already on the CPU.

    x_col = x[1]
    forcing_nt = map(identity, x[2])
    targets_nt = map(identity, y[1])
    masks_nt = map(identity, y[2])

The reason behind is is if the inputs are DimArrays they need to be collected explicitly, thus the Array calls

The implementation of R² using logical indexing (e.g., y[y_nan]) is inefficient on GPUs because it often triggers scalar indexing, which is extremely slow. A more GPU-friendly approach is to use the mask to zero out invalid entries and then perform vectorized reductions. Also, note that this change shifts the definition of R² from squared Pearson correlation to the coefficient of determination, which might be a breaking change for users.

function loss_fn(ŷ, y, y_nan, ::Val{:r2})
    ss_res = sum(abs2, (y .- ŷ) .* y_nan)
    y_mean = sum(y .* y_nan) / sum(y_nan)
    ss_tot = sum(abs2, (y .- y_mean) .* y_nan)
    return 1 - ss_res / ss_tot
end

The use of _get_target_y for a variable named y_nan is semantically confusing. It is better to use _get_target_nan, which was specifically defined for this purpose. Also, please remove the commented-out code to keep the codebase clean.

                y_nan_t = _get_target_nan(y_nan, target)
                _apply_loss(ŷ_t, y_t, y_nan_t, loss_spec)